package health_system.health_system.security;

import health_system.health_system.config.AppProperties;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.JwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.io.Decoders;
import io.jsonwebtoken.security.Keys;
import jakarta.annotation.PostConstruct;
import jakarta.servlet.http.Cookie;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.core.userdetails.User;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.crypto.SecretKey;
import java.util.Arrays;
import java.util.Collection;
import java.util.Date;
import java.util.stream.Collectors;

@Slf4j
@Component
@RequiredArgsConstructor
public class JwtTokenProvider {

    private final AppProperties appProperties;
    private SecretKey secretKey;
    
    private static final String JWT_COOKIE_NAME = "jwt_token";
    
    @PostConstruct
    protected void init() {
        String secret = appProperties.getJwt().getSecret();
        secretKey = Keys.hmacShaKeyFor(Decoders.BASE64.decode(secret));
    }
    
    public String createToken(Authentication authentication) {
        String username = authentication.getName();
        Collection<? extends GrantedAuthority> authorities = authentication.getAuthorities();
        String authoritiesString = authorities.stream()
                .map(GrantedAuthority::getAuthority)
                .collect(Collectors.joining(","));
                
        Date now = new Date();
        Date expiry = new Date(now.getTime() + appProperties.getJwt().getExpiration());
        
        return Jwts.builder()
                .setSubject(username)
                .claim("authorities", authoritiesString)
                .setIssuedAt(now)
                .setExpiration(expiry)
                .signWith(secretKey)
                .compact();
    }
    
    public String createTokenForUser(String username, String authorities) {
        Date now = new Date();
        Date expiry = new Date(now.getTime() + appProperties.getJwt().getExpiration());
        
        return Jwts.builder()
                .setSubject(username)
                .claim("authorities", authorities)
                .setIssuedAt(now)
                .setExpiration(expiry)
                .signWith(secretKey)
                .compact();
    }
    
    public Authentication getAuthentication(String token) {
        Claims claims = Jwts.parserBuilder()
                .setSigningKey(secretKey)
                .build()
                .parseClaimsJws(token)
                .getBody();
                
        String username = claims.getSubject();
        String authoritiesStr = claims.get("authorities", String.class);
        
        Collection<? extends GrantedAuthority> authorities = Arrays.stream(authoritiesStr.split(","))
                .map(SimpleGrantedAuthority::new)
                .collect(Collectors.toList());
                
        User principal = new User(username, "", authorities);
        
        return new UsernamePasswordAuthenticationToken(principal, token, authorities);
    }
    
    public boolean validateToken(String token) {
        try {
            Jwts.parserBuilder().setSigningKey(secretKey).build().parseClaimsJws(token);
            return true;
        } catch (JwtException | IllegalArgumentException e) {
            log.error("Invalid JWT token: {}", e.getMessage());
        }
        return false;
    }
    
    public String resolveToken(HttpServletRequest request) {
        // First, try to get token from header
        String bearerToken = request.getHeader(appProperties.getJwt().getHeader());
        if (StringUtils.hasText(bearerToken) && bearerToken.startsWith(appProperties.getJwt().getPrefix())) {
            return bearerToken.substring(appProperties.getJwt().getPrefix().length() + 1);
        }
        
        // If not found in header, try to get from cookie
        Cookie[] cookies = request.getCookies();
        if (cookies != null) {
            for (Cookie cookie : cookies) {
                if (JWT_COOKIE_NAME.equals(cookie.getName())) {
                    log.debug("Found JWT token in cookie");
                    return cookie.getValue();
                }
            }
        }
        
        return null;
    }
    
    /**
     * Add JWT token to response as a cookie
     * 
     * @param response The HTTP response
     * @param token The JWT token
     */
    public void addCookieToResponse(HttpServletResponse response, String token) {
        int maxAge = (int) (appProperties.getJwt().getExpiration() / 1000);
        
        Cookie cookie = new Cookie(JWT_COOKIE_NAME, token);
        cookie.setPath("/");
        cookie.setHttpOnly(true);
        
        // In production, this should be true, but for development environment
        // we can set it to false if not using HTTPS
        boolean isSecure = appProperties.getJwt().isSecureCookie();
        cookie.setSecure(isSecure);
        
        cookie.setMaxAge(maxAge);
        response.addCookie(cookie);
        
        log.debug("Added JWT token to response cookie with secure={}", isSecure);
    }
    
    /**
     * Delete JWT cookie from response
     * 
     * @param response The HTTP response
     */
    public void deleteCookieFromResponse(HttpServletResponse response) {
        Cookie cookie = new Cookie(JWT_COOKIE_NAME, null);
        cookie.setPath("/");
        cookie.setHttpOnly(true);
        
        // Match the same setting as for adding cookie
        boolean isSecure = appProperties.getJwt().isSecureCookie();
        cookie.setSecure(isSecure);
        
        cookie.setMaxAge(0);
        response.addCookie(cookie);
        
        log.debug("Deleted JWT token cookie");
    }
    
    /**
     * Get the access token expiration time in milliseconds
     * 
     * @return the expiration time in milliseconds
     */
    public Long getAccessTokenExpiration() {
        return appProperties.getJwt().getExpiration();
    }
} 